package com.ms.websocket;

import jakarta.websocket.*;
import jakarta.websocket.server.PathParam;
import jakarta.websocket.server.ServerEndpoint;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.util.Collections;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.CopyOnWriteArraySet;

@Component
@ServerEndpoint("/ws/{sid}")
@Slf4j
public class WebSocketServer {
    // 使用线程安全的连接管理 (支持用户多设备登录)
    private static final ConcurrentMap<Integer, Set<Session>> userSessionMap = new ConcurrentHashMap<>();

    /**
     * 连接建立成功
     */
    @OnOpen
    public void onOpen(Session session, @PathParam("sid") String rawSid) {
        try {
            // 解析用户ID (格式: userId_token)
            int userId = Integer.parseInt(rawSid.split("_")[0]);

            // 添加会话到用户映射
            userSessionMap
                    .computeIfAbsent(userId, k -> new CopyOnWriteArraySet<>())
                    .add(session);

            log.info("用户[{}]建立WebSocket连接, sessionId={}", userId, session.getId());
        } catch (Exception e) {
            log.error("非法连接: sid={}", rawSid, e);
            closeSessionSafely(session);
        }
    }

    /**
     * 定向消息推送
     */
    public static void sendToUser(Integer userId, String message) {
        if (userId == null || message == null) return;

        Set<Session> sessions = userSessionMap.getOrDefault(userId, Collections.emptySet());

        // 遍历用户所有会话
        sessions.removeIf(session -> !session.isOpen()); // 清理无效连接

        sessions.forEach(session -> {
            try {
                // 发送消息
                session.getBasicRemote().sendText(message);
            } catch (IOException e) {
                log.error("消息推送失败 userId={}, session={}", userId, session.getId(), e);
            }
        });
    }

    /**
     * 接收客户端消息
     */
    @OnMessage
    public void onMessage(String message, Session session, @PathParam("sid") String sid) {
        log.info("收到客户端 {} 消息: {}", sid, message);

        // 示例：心跳响应
        if ("PING".equalsIgnoreCase(message.trim())) {
            try {
                // 修复：使用传参的session变量
                session.getBasicRemote().sendText("PONG");
            } catch (IOException e) {
                log.error("心跳响应失败", e);
            }
        }
    }

    /**
     * 连接关闭
     */
    @OnClose
    public void onClose(Session session, @PathParam("sid") String sid) {
        try {
            // 解析用户ID
            int userId = Integer.parseInt(sid.split("_")[0]);

            // 从映射中移除会话
            Set<Session> sessions = userSessionMap.get(userId);
            if (sessions != null) {
                sessions.remove(session);

                // 清理空用户
                if (sessions.isEmpty()) {
                    userSessionMap.remove(userId);
                }
            }

            log.info("连接关闭: userId={}, sessionId={}", userId, session.getId());
        } catch (Exception e) {
            log.error("连接关闭错误: sid={}", sid, e);
        }
    }

    /**
     * 安全关闭会话
     */
    private void closeSessionSafely(Session session) {
        try {
            if (session != null && session.isOpen()) {
                session.close(new CloseReason(CloseReason.CloseCodes.VIOLATED_POLICY, "Invalid session ID"));
            }
        } catch (IOException e) {
            log.error("关闭会话失败", e);
        }
    }
}